class Solution(object):
    def countSquares(self, matrix):
        m = len(matrix)
        n = len(matrix[0])
        ans = 0
        for i in range(1, m):
            matrix[i][0] += matrix[i - 1][0]
        for j in range(1, n):
            matrix[0][j] += matrix[0][j - 1]
        for i in range(1, m):
            for j in range(1, n):
                matrix[i][j] += matrix[i][j - 1] + matrix[i - 1][j] - matrix[i - 1][j - 1]
        for i in range(m):
            for j in range(n):
                for k in range(min(m, n)):
                    if i + k < m and j + k < n:
                        t = matrix[i + k][j + k]
                        if i > 0:
                            t -= matrix[i - 1][j + k]
                        if j > 0:
                            t -= matrix[i + k][j - 1]
                        if i > 0 and j > 0:
                            t += matrix[i - 1][j - 1]
                        if t == (k + 1) ** 2:
                            ans += 1
                        else:
                            break
                    else:
                        break
        return ans


data = Solution()
matrix = [
    [0, 1, 1, 1],
    [1, 1, 1, 1],
    [0, 1, 1, 1]
]
print(data.countSquares(matrix))

data = Solution()
matrix = [
    [1, 0, 1],
    [1, 1, 0],
    [1, 1, 0]
]
print(data.countSquares(matrix))
